/**
 * Copyright (C) 2025 pony working team
 * SPDX-License-Identifier: MulanPSL-2.0
 */

#include <gtest/gtest.h>

#include <stdlib.h>

#include <x_common.h>
#include <x_rbtree.h>

extern "C" {

typedef struct item {
    rb_node_t node;
    int k;
    int v;
} item_t;

typedef struct map {
    rb_tree_t items;
} map_t;

static int map_compare_nk(bi_node_t *n, uintptr_t k)
{
    item_t *i = container_of(n, item_t, node.bin);

    if ((uintptr_t)i->k == k) {
        return 0;
    }

    return (uintptr_t)i->k > k ? 1 : -1;
}

static int map_compare_nn(bi_node_t *a, bi_node_t *b)
{
    return map_compare_nk(a, (uintptr_t)container_of(b, item_t, node.bin)->k);
}

int map_init(map_t *m)
{
    rb_init(&m->items, map_compare_nn, map_compare_nk);

    return 0;
}

int map_insert(map_t *m, int k, int v)
{
    item_t *i = (item_t *)malloc(sizeof(item_t));
    i->k = k;
    i->v = v;

    return rb_insert(&m->items, &i->node);
}

int map_value(map_t *m, int k)
{
    rb_node_t *rbn = rb_find_eq(&m->items, k);
    if (rbn == NULL) {
        return -1;
    }
    item_t *i = container_of(rbn, item_t, node);
    return i->v;
}

void map_erase(map_t *m, int k)
{
    rb_node_t *rbn = rb_find_eq(&m->items, k);
    if (rbn == NULL) {
        return;
    }
    rb_delete(&m->items, rbn);
    item_t *i = container_of(rbn, item_t, node);
    free(i);
}

int map_foreach_cb(bi_node_t *n)
{
    rb_node_t *l = NULL;
    rb_node_t *p = NULL;
    rb_node_t *r = NULL;
    rb_node_t *w = container_of(n, rb_node_t, bin);

    if (n->l) {
        l = container_of(n->l, rb_node_t, bin);
    }
    if (n->p) {
        p = container_of(n->p, rb_node_t, bin);
    }
    if (n->r) {
        r = container_of(n->r, rb_node_t, bin);
    }

    item_t *i = container_of(w, item_t, node);
    printf("k=%d,v=%d,c=%d", i->k, i->v, w->c);
    if (l) {
        i = container_of(l, item_t, node);
        printf(",lk=%d", i->k);
    } else {
        printf(",lk=x");
    }
    if (p) {
        i = container_of(p, item_t, node);
        printf(",pk=%d", i->k);
    } else {
        printf(",pk=x");
    }
    if (r) {
        i = container_of(r, item_t, node);
        printf(",rk=%d", i->k);
    } else {
        printf(",rk=x");
    }
    printf("\n");

    return 0;
}

void map_dump(map_t *m)
{
    rb_foreach(&m->items, map_foreach_cb);
    printf("\n");
}
}

class TestRbtree : public ::testing::Test
{
protected:
    void SetUp() {};

    void TearDown() {};
};

TEST_F(TestRbtree, map)
{
    map_t m;
    EXPECT_EQ(map_init(&m), 0);

    /**
     * b - 1
     * r - 0
     */

    EXPECT_EQ(map_insert(&m, 1, 1), 0);
    /**
     *     1 (b)
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 5, 5), 0);
    /**
     *     1 (b)
     *      \
     *       5 (r)
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 4, 4), 0);
    /**
     *         4 (b)
     *        / \
     *   (r) 1   5 (r)
     *
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 3, 3), 0);
    /**
     *         4 (b)
     *        / \
     *   (b) 1   5 (b)
     *        \
     *         3 (r)
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 8, 8), 0);
    /**
     *          4 (b)
     *        /   \
     *  (b) 1       5 (b)
     *       \       \
     *        3 (r)   8 (r)
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 6, 6), 0);
    /**
     *           4 (b)
     *        /     \
     *  (b) 1          6 (b)
     *       \      /    \
     *    (r) 3    5 (r)  8 (r)
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 2, 2), 0);
    /**
     *                 4 (b)
     *              /     \
     *       (b) 2           6 (b)
     *         /   \      /    \
     *   (r) 1   (r) 3   5 (r)  8 (r)
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 7, 7), 0);
    /**
     *                 4 (b)
     *              /     \
     *       (b) 2           6 (r)
     *         /   \      /    \
     *   (r) 1   (r) 3   5 (b)  8 (b)
     *                        /
     *                     7 (r)
     */
    map_dump(&m);

    EXPECT_EQ(map_insert(&m, 9, 9), 0);
    /**
     *                 4 (b)
     *              /     \
     *       (b) 2           6 (r)
     *         /   \      /    \
     *   (r) 1   (r) 3   5 (b)  8 (b)
     *                        /   \
     *                     7 (r)   9 (r)
     */
    map_dump(&m);

    EXPECT_EQ(map_value(&m, 1), 1);
    EXPECT_EQ(map_value(&m, 2), 2);
    EXPECT_EQ(map_value(&m, 3), 3);
    EXPECT_EQ(map_value(&m, 4), 4);

    map_erase(&m, 1);
    EXPECT_EQ(map_value(&m, 1), -1);
    /**
     *                 4 (b)
     *              /     \
     *       (b) 2           6 (r)
     *             \      /    \
     *          (r) 3   5 (b)  8 (b)
     *                        /   \
     *                     7 (r)   9 (r)
     */
    map_dump(&m);

    map_erase(&m, 8);
    /**
     *                 4 (b)
     *              /     \
     *       (b) 2           6 (r)
     *             \      /    \
     *          (r) 3   5 (b)  9 (b)
     *                        /
     *                     7 (r)
     */
    map_dump(&m);

    map_erase(&m, 2);
    /**
     *             4 (b)
     *            / \
     *       (b) 3   6 (r)
     *              / \
     *         (b) 5   9 (b)
     *                /
     *               7 (r)
     */
    map_dump(&m);

    map_erase(&m, 4);
    /**
     *             5 (b)
     *            / \
     *       (b) 3   7 (r)
     *              / \
     *         (b) 6   9 (b)
     */
    map_dump(&m);

    map_erase(&m, 7);
    /**
     *             5 (b)
     *            / \
     *       (b) 3   9 (b)
     *              /
     *         (r) 6
     */
    map_dump(&m);

    map_erase(&m, 3);
    /**
     *             6 (b)
     *            / \
     *       (b) 5   9 (b)
     */
    map_dump(&m);

    map_erase(&m, 6);
    /**
     *             9 (b)
     *            /
     *       (r) 5
     */
    map_dump(&m);

    map_erase(&m, 9);
    /**
     *             5 (b)
     */
    map_dump(&m);

    map_erase(&m, 5);
}
